import cv2
import matplotlib.pyplot as plt
import sys
import numpy as np


def padding_img(img):
    """
    padding image
    :param img: (height, width, channel)
    """
    height, width, channels = img.shape

    padding_size_1 = abs(height - width) // 2
    padding_size_2 = abs(height - width) - padding_size_1

    if padding_size_1 == 0:
        img_out = img
    else:
        # padding width axis
        if height > width:
            padding1 = np.zeros(shape=(height, padding_size_1, channels), dtype=np.uint8)
            padding2 = np.zeros(shape=(height, padding_size_2, channels), dtype=np.uint8)
            img_out = np.concatenate([padding1, img, padding2], axis=1)
            # print(img_out)
        # padding height axis
        elif height < width:
            padding1 = np.zeros(shape=(padding_size_1, width, channels), dtype=np.uint8)
            padding2 = np.zeros(shape=(padding_size_2, width, channels), dtype=np.uint8)
            img_out = np.concatenate([padding1, img, padding2], axis=0)
        else:
            img_out = img

    return img_out


def read_img(img_file, resize=(224, 224)):
    img = cv2.imread(img_file, 1)
    # print(img_file)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    print(img)
    print(img.shape)
    img = padding_img(img)
    print(img.shape)
    img = cv2.resize(img, resize)
    # print(img.shape)
    return img  # .reshape((3, resize[0], resize[1]))


if __name__ == '__main__':
    infile = sys.argv[1]
    img = read_img(infile)
    print(img.shape)
    plt.imshow(img)
    plt.show()
